///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/gradient_acc.h
///
/// \brief   header file for the gradient accumulator class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef FORMIC_FQMC_GRADIENT_ACC_HEADER
#define FORMIC_FQMC_GRADIENT_ACC_HEADER

#include<vector>
#include<string>
#include<numeric>
#include<cassert>
#include<algorithm>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/fqmc/accumulator.h>
#include<formic/wfn/configuration.h>
#include<formic/fqmc/vmc_walker.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/output/output.h>
#include<formic/exception.h>
#include<formic/lapack/interface.h>

namespace formic { namespace fqmc {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  a class to accumulate the gradient during a markov chain evolution
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template <class S> class GradAcc : public formic::Accumulator<S> {

    // data members
    private:

      /// \brief  pointer to the guiding function
      boost::shared_ptr<formic::CoeffObj<S> > _gf;

      /// \brief  pointer to the value function
      boost::shared_ptr<formic::CoeffObj<S> > _vf;

//      /// \brief  reference to the walker's configuration
//      const formic::Configuration & _config;

      /// \brief  energy
      double _grad_norm;

      /// \brief  energy
      double _energy;

      /// \brief  energy variance
      double _variance;

      /// \brief  weight accumulated on local process
      double _total_proc_weight;

      /// \brief  pointer to the energy (null pointer if energy not given)
      const double * _energy_ptr;

      /// \brief  a history of the sampled probabilities
      std::vector<double> _history_p;

      /// \brief  a history of the |value/guiding|^2 ratios
      std::vector<double> _history_v;

      /// \brief  a history of the sampled local energies
      std::vector<double> _history_e;

      /// \brief  whether to ignore (i.e. skip) derivatives for fixed variables
      bool _ign_fixed;

      /// \brief  the <dPsi/dx_i|H|Psi> vector
      std::vector<S> _dhvec;

      /// \brief  the <dPsi/dx_i|Psi> vector
      std::vector<S> _dpvec;

      /// \brief  vector to hold the final gradient in compact (i.e. complex number) form
      std::vector<S> _cgrad;

      /// \brief  vector of large weights
      std::vector<double> _large_weights;

      /// \brief  vector of configurations with large weights
      std::vector<size_t> _large_configs;

    // private member functions
    private:

      void negate_imaginary_part(const int n, const double * const c, double * const r) const {}
      void negate_imaginary_part(const int n, const std::complex<double> * const c, double * const r) const {
        for (int i = 0; i < n; i++)
          r[i+n] = -r[i+n];
      }

    public:

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  constructs the accumulator around a specific coefficient object
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      GradAcc(const boost::shared_ptr<formic::CoeffObj<S> > & gf,
              const boost::shared_ptr<formic::CoeffObj<S> > & vf,
              //const formic::Configuration & config,
              const bool ign_fixed = true,
              const double * const energy_ptr = 0)
        : _gf(gf),
          _vf(vf),
          //_config(config),
          _grad_norm(0.0),
          _energy(0.0),
          _variance(0.0),
          _total_proc_weight(0.0),
          _energy_ptr(energy_ptr),
          _history_p(),
          _history_v(),
          _history_e(),
          _ign_fixed(ign_fixed),
          _dhvec( ( ign_fixed ? vf->n_der() : vf->n_var() ), formic::zero(S())),
          _dpvec( ( ign_fixed ? vf->n_der() : vf->n_var() ), formic::zero(S())),
          _cgrad( ( ign_fixed ? vf->n_der() : vf->n_var() ), formic::zero(S())),
          _large_weights(),
          _large_configs()
      {}

      /// \brief  return the accumulator's name
      std::string name() const { return "gradient accumulator"; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  Prints some statistics about the most recent sample.
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void print_statistics(formic::Output<std::ofstream> & fout) {

        if ( formic::mpi::rank() == 0 ) {
          fout << boost::format("gradient accumulation statistics on process %i:") % formic::mpi::rank() << std::endl;
          fout << std::endl;
          fout << boost::format("   acceptance ratio = %20.12f") % (this->_naccept / this->_total_weight) << std::endl;
          fout << boost::format("      gradient norm = %20.12f") % _grad_norm << std::endl;
          fout << boost::format("             energy = %20.12f") % _energy << std::endl;
          fout << boost::format("    energy variance = %20.12f") % _variance << std::endl;
          fout << boost::format("   energy std. dev. = %20.12f") % std::sqrt(_variance) << std::endl;
          fout << std::endl;
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Accumulate the local energy and gradient contributions
      ///
      /// \param[in]      userinp  the user's input options
      /// \param[in]      w        the weight to apply to this configuration's contribution
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_accumulate(const formic::InputBase & userinp, const double w) {

        // record weight for local process
        _total_proc_weight += w;

        // update the weight histories
        _history_p.push_back(w);

        // get the square of the value/guiding ratio
        const double vg_square = ( _vf == _gf  ?  1.0  :  formic::square_norm( ( _vf->get_coeff() / _gf->get_coeff() ).get() ) );

        // update the value/guiding ratio history
        _history_v.push_back(vg_square);

        // get the local energy
        const S local_energy = _vf->local_energy();

        // update the energy history
        _history_e.push_back( formic::real(local_energy) );

        // get the derivative ratios
        _vf->der_ratios(&_cgrad.at(0), _ign_fixed);

        // if we have the total energy, add the local gradent as one vector
        if ( _energy_ptr != 0 ) {

          const S energy_diff = local_energy - formic::unity(S()) * (*_energy_ptr);
          formic::xaxpy(_cgrad.size(), w * vg_square * energy_diff, &_cgrad.at(0), 1, &_dhvec.at(0), 1);

        // if we don't have the total energy, then save the gradient pieces separately
        } else {

          // update the <dPsi|H|Psi> vector
          formic::xaxpy(_cgrad.size(), w * vg_square *       local_energy, &_cgrad.at(0), 1, &_dhvec.at(0), 1);

          // update the <dPsi|Psi> vector
          formic::xaxpy(_cgrad.size(), w * vg_square * formic::unity(S()), &_cgrad.at(0), 1, &_dpvec.at(0), 1);

        }

        // record configurations with the largest weights
        if ( formic::mpi::rank() == 0 && userinp.get<int>("configs_to_print") > 0 ) {
          throw formic::Exception("printing large configurations is no longer implemented");
          //size_t cmpd = 0;
          //size_t i;
          //for (i = 0; i < _config.fock_config().size(); i++)
          //  cmpd = size_t(2)*cmpd + size_t(_config.fock_config()[i]);
          //for (i = 0; i < _large_configs.size(); i++) {
          //  if ( _large_configs[i] == cmpd )
          //    break;
          //}
          //if ( i < _large_configs.size() ) {
          //  //formic::of << "adding to existing weight" << std::endl;
          //  _large_weights[i] += w;
          ////} else if ( _large_configs.size() < userinp.get<int>("configs_to_print") ) {
          ////  //formic::of << "adding new weight" << std::endl;
          ////  _large_configs.push_back(cmpd);
          ////  _large_weights.push_back(w);
          ////  i = _large_configs.size() - 1;
          //} else {
          //  bool not_placed = true;
          //  for (i = 0; i < _large_weights.size(); i++) {
          //    if ( _large_weights[i] < w ) {
          //      //formic::of << "inserting new weight" << std::endl;
          //      _large_configs.insert(_large_configs.begin()+i, cmpd);
          //      _large_weights.insert(_large_weights.begin()+i, w);
          //      //_large_configs.pop_back();
          //      //_large_weights.pop_back();
          //      not_placed = false;
          //      break;
          //    }
          //  }
          //  if ( not_placed ) {
          //    _large_configs.push_back(cmpd);
          //    _large_weights.push_back(w);
          //  }
          //  i = 0;
          //}
          ////formic::of << "stored configs (unsorted)";
          ////for (size_t k = 0; k < _large_configs.size(); k++)
          ////  formic::of << boost::format(" %16i") % _large_configs.at(k);
          ////formic::of << std::endl;
          ////formic::of << "stored weights (unsorted)";
          ////for (size_t k = 0; k < _large_weights.size(); k++)
          ////  formic::of << boost::format(" %16.6e") % _large_weights.at(k);
          ////formic::of << std::endl;
          //if (i > 0) {
          //  size_t j;
          //  for (j = 0; j < i && _large_weights[j] > _large_weights[i]; j++) {}
          //  //formic::of << "j = " << j << std::endl;
          //  for (size_t k = i; k > j; k--) {
          //    std::swap(_large_weights[k], _large_weights[k-1]);
          //    std::swap(_large_configs[k], _large_configs[k-1]);
          //  }
          //}
          ////formic::of << "stored configs (  sorted)";
          ////for (size_t k = 0; k < _large_configs.size(); k++)
          ////  formic::of << boost::format(" %16i") % _large_configs.at(k);
          ////formic::of << std::endl;
          ////formic::of << "stored weights (  sorted)";
          ////for (size_t k = 0; k < _large_weights.size(); k++)
          ////  formic::of << boost::format(" %16.6e") % _large_weights.at(k);
          ////formic::of << std::endl;
          ////formic::of << std::endl;
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief finalizes the accumulation (this may include for example communicating the results)
      ///
      /// \param[in]      userinp  the user's input options
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_finalize(const formic::InputBase & userinp) {

        // normalize the probabilities
        formic::xscal(_history_p.size(), 1.0 / this->_total_weight, &_history_p.at(0), 1);

        // compute average of value/guiding ratios
        double vga = 0.0;
        {
          double x = 0.0;
          for (int i = 0; i < _history_v.size(); i++)
            x += _history_v.at(i) * _history_p.at(i);
          formic::mpi::allreduce(&x, &vga, 1, MPI::SUM);
        }

        // update local energies using value/guiding ratio weights
        for (int i = 0; i < _history_v.size(); i++)
          _history_e.at(i) *= _history_v.at(i) / vga;

        // compute average of local energies
        _energy = 0.0;
        {
          double x = 0.0;
          for (int i = 0; i < _history_e.size(); i++)
            x += _history_e.at(i) * _history_p.at(i);
          formic::mpi::allreduce(&x, &_energy, 1, MPI::SUM);
        }

        // compute the variance
        _variance = 0.0;
        {
          double var = 0.0;
          for (int i = 0; i < _history_p.size(); i++) {
            assert( _history_p.at(i) >= 0.0 );
            var += _history_p.at(i) * formic::square_norm(_energy - _history_e.at(i));
          }
          formic::mpi::allreduce(&var, &_variance, 1, MPI::SUM);
        }

        // compute the gradient contributions and collect them on the root process
        if ( _energy_ptr == 0 )
          formic::xaxpy(_cgrad.size(), -formic::unity(S()) * _energy, &_dpvec.at(0), 1, &_dhvec.at(0), 1);
        formic::xscal(_cgrad.size(), 2.0 * formic::unity(S()) / ( this->_total_weight * vga ), &_dhvec.at(0), 1);
        formic::mpi::reduce(&_dhvec.at(0), &_cgrad.at(0), _cgrad.size(), MPI::SUM);

        // compute the gradient norm and send it to all processes
        if ( formic::mpi::rank() == 0 )
          _grad_norm = std::abs( std::sqrt( formic::xdotc(_cgrad.size(), &_cgrad.at(0), 1, &_cgrad.at(0), 1) ) );
        formic::mpi::bcast(_grad_norm);

        // print configurations with the largest weights
        if ( formic::mpi::rank() == 0 && userinp.get<int>("configs_to_print") > 0 ) {
          throw formic::Exception("printing large configurations is no longer implemented");
          //for (size_t i = 0; i < _large_weights.size(); i++)
          //  _large_weights.at(i) /= _total_proc_weight;
          //formic::of << boost::format("printing largest %i configurations (out of %i total visited)")
          //                % userinp.get<int>("configs_to_print")
          //                % _large_weights.size()
          //           << std::endl
          //           << std::endl;
          //for (size_t i = 0; i < _large_weights.size() && i < userinp.get<int>("configs_to_print"); i++) {
          //  std::vector<size_t> cfg(_config.fock_config().size(), 0);
          //  size_t cmpd = _large_configs.at(i);
          //  for (int j = cfg.size() - 1; j >= 0; j--) {
          //    cfg.at(j) = cmpd % 2;
          //    cmpd /= 2;
          //  }
          //  for (size_t j = 0; j+1 < cfg.size(); j += 2)
          //    formic::of << "  " << cfg.at(j) << cfg.at(j+1);
          //  formic::of << boost::format("      %10.4e") % _large_weights.at(i) << std::endl;
          //}
          //formic::of << std::endl;
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  resets the accumulator
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_reset(const formic::Configuration & config) {
        _grad_norm = 0.0;
        _energy = 0.0;
        _variance = 0.0;
        _total_proc_weight = 0.0;
        _history_p.clear();
        _history_v.clear();
        _history_e.clear();
        std::fill(_dhvec.begin(), _dhvec.end(), formic::zero(S()));
        std::fill(_dpvec.begin(), _dpvec.end(), formic::zero(S()));
        _large_weights.clear();
        _large_configs.clear();
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  return the energy
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      double energy() const { return _energy; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  return the variance
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      double variance() const { return _variance; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  return the gradient norm
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      double grad_norm() const { return _grad_norm; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  Get the gradient in real number format.  For complex wave functions, this means that
      ///         the gradient is returned as the energy derivatives with respect to the varaibles' real
      ///         parts followed by the energy derivatives with respect to the variables' imaginary parts
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void get_real_gradient(double * const grad) const {
        formic::expand_to_real(_cgrad.size(), &_cgrad.at(0), grad);
        this->negate_imaginary_part(_cgrad.size(), &_cgrad.at(0), grad);
      }

  };

} // end namespace fqmc

} // end namespace formic

#endif
